Move Gemma4 per-layer embeddings to embedding model (Alternative A) - #296

…318)

## Problem

L4 / L5 e2e tests have been failing on `main` for every Gemma4 case
since
#296 (Move Gemma4 per-layer embeddings to embedding model). Two distinct
root causes:

### 1. `per_layer_inputs` missing from decoder feed

After #296 the embedding sub-model emits a second output
`per_layer_inputs` and the decoder accepts it as a required input. The
e2e harness only wired the first embedding output (`inputs_embeds`)
into the decoder, so every multi-model Gemma4 path failed with:

```
ValueError: Required inputs (['per_layer_inputs']) are missing from input feed
(['inputs_embeds', 'past_key_values.0.key', ..., 'attention_mask', 'position_ids'])
```

Affected: text-only on multi-model (`text-generation/gemma-4-e2b`),
VL prefill (`image-text-to-text/gemma-4-e2b-it`, `…-e4b-it`), VL
generation (L5), speech-language prefill+generation.

### 2. `input_features_mask` dtype mismatch

The Gemma4 audio encoder (`_gemma4.py:367`) declares
`input_features_mask` as `tensor(bool)`, but the harness unconditionally
cast every feature-extractor output to `np.float32`, producing:

```
InvalidArgument: Unexpected input data type. Actual: (tensor(float)), expected: (tensor(bool))
```

Even when the fallback all-True mask path was taken (line 1267), the
bool numpy array crashed `ort_easy`'s DLPack-first conversion path
because DLPack has no native bool type code.

## Fix

### `tests/e2e_golden_test.py`

In every decoder-feed setup (5 sites: text-only multi-model prefill,
VL prefill, VL generation, speech-language prefill, speech-language
generation), wire any extra embedding outputs through to the decoder by
name. For models without `per_layer_inputs` the extra loop iteration is
a no-op:

```python
elif name in emb_out:
    dec_feeds[name] = emb_out[name]
```

For both audio-encoder setup sites (L4 + L5 paths), honor the session's
declared input dtype:

```python
target_dtype = audio_session.get_input_dtype(name) or np.float32
audio_feeds[name] = audio_processed[name].astype(target_dtype)
```

And use the session-declared dtype for the fallback all-True mask too.

### `src/mobius/_testing/ort_inference.py`

Route bool numpy arrays through `OrtValue.ortvalue_from_numpy` directly
in `_numpy_to_ort_value`, bypassing `ort_easy`'s DLPack-first path
(which has no bool type code).

## Verification (locally on H200)

| Test | Before | After |
|---|---|---|
| L4 `text-generation/gemma-4-e2b` | FAIL (missing per_layer_inputs) |
**PASS** |
| L4 `image-text-to-text/gemma-4-e2b-it` | FAIL (missing
per_layer_inputs) | **PASS** |
| L4 `speech-language/gemma-4-e2b-it-audio` | FAIL (bool dtype mismatch)
| **PASS** |
| L5 `image-text-to-text/gemma-4-e2b-it` | FAIL | **PASS** |
| L5 `speech-language/gemma-4-e2b-it-audio` | FAIL | **PASS** |

5 passed, 1 skipped (L5 text-gen is gated by integration markers under
fast runs), 0 failed. `ruff check` + `ruff format --check` both pass.

L1 + L3 + non-gemma e2e suites are unaffected (the embedding-output
loop is a generic no-op for models that don't emit extra outputs;
audio-mask changes only trigger when the session declares a bool input).

## Why this didn't get caught earlier

The CI matrix only runs L4 / L5 in the affected-models lane, and #296
was tested standalone before this lane started exercising the multi-
model Gemma4 path through the e2e harness. The bool-mask issue is even
older — it was masked by the harness's unconditional
`astype(np.float32)`
until the audio_encoder was upgraded to take a real BOOL mask.

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

…o decoder

After #296 ("Move Gemma4 per-layer embeddings to embedding model") the
Gemma4 embedding sub-model emits a second output ``per_layer_inputs``
([B, S, num_layers * per_layer_dim]) that the decoder consumes on every
step (prefill + decode). The example wasn't updated to forward this
output, so any Gemma4 build with hidden_size_per_layer_input > 0
(e.g. google/gemma-4-E2B-it) blows up at the first decoder.run() with:

  ValueError: Required inputs (['per_layer_inputs']) are missing from
  input feed (['inputs_embeds', ..., 'attention_mask', 'position_ids'])

Fix:

  * prepare_decoder_feeds(): accept an optional per_layer_inputs
    argument and add it to the feeds dict when present. Builds without
    per-layer inputs (hidden_size_per_layer_input == 0, e.g. larger
    Gemma 4 variants) keep working since the kwarg is optional.

  * generate(): pull embed_out.get("per_layer_inputs") and pass it to
    prepare_decoder_feeds on every step.

Verified locally: 'python examples/gemma4_multimodal.py --mode text
--prompt "What is 2+2?"' now generates "2 + 2 = **4**".

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

## Summary

After #296 ("Move Gemma4 per-layer embeddings to embedding model") the
Gemma4 embedding sub-model emits a second output `per_layer_inputs`
(`[B, S, num_layers * per_layer_dim]`) that the decoder consumes on
every step. `examples/gemma4_multimodal.py` wasn't updated to forward
this output, so any Gemma 4 build with `hidden_size_per_layer_input > 0`
(e.g. `google/gemma-4-E2B-it`) crashes at the first `decoder.run()`:

```
ValueError: Required inputs (['per_layer_inputs']) are missing from input feed
(['inputs_embeds', ..., 'attention_mask', 'position_ids'])
```

## Fix

- `prepare_decoder_feeds()`: accept an optional `per_layer_inputs`
argument and add it to the feeds dict when present. Builds without
per-layer inputs (`hidden_size_per_layer_input == 0`, e.g. larger Gemma
4 variants) keep working — the kwarg is optional.
- `generate()`: pull `embed_out.get("per_layer_inputs")` and pass it to
`prepare_decoder_feeds` on every step.

## Verification

```
$ python examples/gemma4_multimodal.py --mode text --prompt "What is 2+2?"
...
📝  TEXT-ONLY GENERATION
================================================================
Prompt: What is 2+2?
----------------------------------------------------------------
2 + 2 = **4**
```

## Related

This is the example-side counterpart to #318 which fixed the same gap in
the L4/L5 e2e test harness.

---------

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Co-authored-by: justinchuby <11205048+justinchuby@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

The internal input_ids path and compute_per_layer_inputs kwarg depended
on gemma4.py changes that were already reverted when we restored the
fused-table architecture from PR #296. Revert to main's external
per_layer_inputs path to keep the two files consistent.

Co-Authored-By: Claude <noreply@anthropic.com>

The internal input_ids path and compute_per_layer_inputs kwarg depended
on gemma4.py changes that were already reverted when we restored the
fused-table architecture from PR #296. Revert to main's external
per_layer_inputs path to keep the two files consistent.

Co-Authored-By: Claude <noreply@anthropic.com>